import matplotlib
matplotlib.use('Agg')  # 或 'TkAgg'，根据运行环境选择
import matplotlib.pyplot as plt
import numpy as np
import os
import cv2

# 数据解析函数
def parse_data(raw_data):
    """
    解析原始数据，将其分组为每个 index 的数据。
    :param raw_data: 原始数据（字符串列表，每行是一个字符串）
    :return: 按 index 分组的字典，每个 index 包含时间戳和点云数据
    """
    parsed_data = {}
    current_index = None

    for line in raw_data:
        line = line.strip()
        if line.startswith("index:"):  # 检测 index 行
            parts = line.split()
            current_index = int(parts[1])  # 提取 index
            timestamp = int(parts[2])  # 提取时间戳
            parsed_data[current_index] = {
                "timestamp": timestamp,
                "points": []
            }
        elif line.startswith("3"):  # 检测点数据
            parts = line.split()

            # 3 254 1.19121 0.272047 -0.0397453
            x, y = map(float, parts[2:4])  # 提取 x, y
            if current_index is not None:
                parsed_data[current_index]["points"].append((x, y))
    return parsed_data

# 坐标映射函数
def map_to_image_coordinates(points, img_size=640, scale=1.5):
    """
    将点的坐标映射到图像像素坐标系。
    :param points: 原始坐标点列表 [(x1, y1), (x2, y2), ...]
    :param img_size: 图像尺寸（正方形，默认为 640x640）
    :param scale: 坐标范围的缩放因子（-scale 到 scale 映射到图像范围）
    :return: 映射后的像素坐标列表 [(px1, py1), (px2, py2), ...]
    """
    center = img_size // 2  # 图像中心点
    pixel_points = []
    for x, y in points:
        px = int(center - 100 * y)  # x 映射到像素
        py = int(center - 100 * x)  # y 映射到像素（注意 y 轴反转）
        pixel_points.append((px, py))
    return pixel_points

# 可视化函数
def visualize_data(parsed_data, img_size=640, output_dir="output"):
    """
    可视化解析后的数据，将点绘制在 640x640 的图像上。
    :param parsed_data: 按 index 分组的字典数据
    :param img_size: 图像尺寸（正方形，默认为 640x640）
    :param output_dir: 保存图像的目录
    """
    os.makedirs(output_dir, exist_ok=True)  # 创建输出目录

    for index, data in parsed_data.items():
        # 创建空白图像
        # img = np.ones((img_size, img_size, 3), dtype=np.uint8) * 255  # 白色背景
        img = np.zeros((img_size, img_size, 3), dtype=np.uint8)  # 黑色背景


        # 绘制红色的原点 (300, 300)
        center_point = (300, 300)  # 原点坐标
        cv2.circle(img, center_point, radius=3, color=(255, 0, 0), thickness=-1)  # 红色点 (BGR)

        # 绘制蓝色矩形 (280, 250) 到 (320, 310)
        top_left_blue = (280, 250)  # 蓝色矩形左上角
        bottom_right_blue = (320, 310)  # 蓝色矩形右下角
        cv2.rectangle(img, top_left_blue, bottom_right_blue, color=(255, 0, 0), thickness=1)  # 蓝色填充矩形 (BGR)

        # 绘制绿色矩形框 (150, 150) 到 (450, 450)
        top_left_green = (150, 150)  # 绿色矩形框左上角
        bottom_right_green = (450, 450)  # 绿色矩形框右下角
        cv2.rectangle(img, top_left_green, bottom_right_green, color=(0, 255, 0), thickness=2)  

        points = data["points"]
        pixel_points = map_to_image_coordinates(points, img_size=img_size)

        # 绘制点
        fig, ax = plt.subplots(figsize=(6, 6))
        ax.imshow(img)  # 显示背景
        for px, py in pixel_points:
            ax.scatter(px, py, c='blue', s=10)  # 绘制点，蓝色，大小为 10

        # 设置标题和坐标轴
        ax.set_title(f"Index: {index}, Timestamp: {data['timestamp']}")
        ax.set_xlim(0, img_size)
        ax.set_ylim(img_size, 0)  # y 轴反转
        ax.axis('off')  # 关闭坐标轴

        # 保存图像
        output_path = os.path.join(output_dir, f"index_{index}.png")
        plt.savefig(output_path)
        plt.close(fig)  # 关闭图像，释放内存



import csv

# 读取 CSV 文件
csv_file = "data.csv"
raw_data = []

with open(csv_file, "r") as file:
    reader = csv.reader(file)
    for row in reader:
        # 每行作为字符串加入 raw_data
        raw_data.append(row[0])

# 解析数据
parsed_data = parse_data(raw_data)

# 可视化数据并保存图像
visualize_data(parsed_data)
